#include "xsocket.h"

#include <unistd.h>
#include <string.h>
#include <fcntl.h>

#include <sys/socket.h>
#include <sys/poll.h>
#include <sys/time.h>
#include <netinet/in.h>
#include <arpa/inet.h>

#include <stdio.h> // for test

uint32_t get_cur_millseconds()
{
	struct timeval now;
	gettimeofday(&now, NULL);
	return now.tv_sec * 1000 + now.tv_usec / 1000;
}


#define SOCK_INIT (-1)
#define IS ==

sock_t xopen(const char * ip_address, int port, SOCK_TYPE st)
{
        sock_t sd = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
        if (sd IS -1) {
                return sd;
        }

        bool ok = false;
        do {
                if (st IS SOCK_CLIENT_TYPE) { /* CLIENT */
                        struct sockaddr_in addr;
                        memset(&addr, 0, sizeof(addr));
                        addr.sin_family = AF_INET;
                        addr.sin_addr.s_addr = inet_addr(ip_address);
                        addr.sin_port = htons(port);

                        if (connect(sd, (const struct sockaddr*)&addr, (socklen_t)sizeof(addr)) IS -1) {
                                break;
                        }

                        int flag = -1;
                        if ((flag=fcntl(sd, F_GETFL)) IS -1
                                        || fcntl(sd, F_SETFL, flag|O_NONBLOCK) IS -1) {
                                break;
                        }

                        struct linger ling;
                        ling.l_onoff = 1;
                        ling.l_linger = 0;
                        if ( setsockopt( sd, SOL_SOCKET, SO_LINGER, &ling, sizeof( ling ) ) IS -1 ) {
                                break;
                        }

                        ok = true;
                } else if (st IS SOCK_SERVER_TYPE) { /* SERVER */
                        /* set reuse */
                        int reuse = 1;
                        if (setsockopt(sd, SOL_SOCKET, SO_REUSEADDR, &reuse, sizeof(reuse)) IS -1) {
                                break;
                        }

                        struct sockaddr_in addr;
                        memset(&addr, 0, sizeof(addr));
                        addr.sin_family = AF_INET;
                        addr.sin_addr.s_addr = inet_addr(ip_address);
                        addr.sin_port = htons(port);

                        if (bind(sd, (const struct sockaddr*)&addr, (socklen_t)sizeof(addr)) IS -1) {
                                break;
                        }

                        if (listen(sd, SOMAXCONN) IS -1) {
                                break;
                        }

                        ok = true;
                }
        } while (false);

        if (!ok) {
                xclose(sd);
                sd = SOCK_INIT;
        }

        return sd;
}

int xclose(sock_t sd)
{
        return close(sd);
}

sock_t xrecv_client(sock_t sd, uint32_t * client_ip, int * client_port, int timeout_ms)
{
        struct pollfd poll_fd;
        memset(&poll_fd, 0, sizeof(poll_fd));

        poll_fd.fd = sd;
        poll_fd.events = (POLLIN | POLLRDNORM);
        poll_fd.revents = 0;
        int poll_ret = poll(&poll_fd, 1, timeout_ms);
        if (poll_ret IS -1) {
                return -1;
        } else if (poll_ret IS 0) {
                return 0;
        }

        if (poll_fd.revents & (POLLERR | POLLHUP | POLLNVAL)) {
                return -2;
        }

        struct sockaddr_in new_addr;
        memset(&new_addr, 0, sizeof(new_addr));
        socklen_t new_addr_len = sizeof(new_addr);

        sock_t new_sd = accept(sd, (struct sockaddr*)&new_addr, &new_addr_len);
        if (new_sd IS -1) {
                return new_sd;
        }

        int flag = -1;
        if ((flag=fcntl(new_sd, F_GETFL)) IS -1
                        || fcntl(new_sd, F_SETFL, flag|O_NONBLOCK) IS -1) {
                xclose(new_sd);
                new_sd = SOCK_INIT;
                return new_sd;
        }

        if (client_ip) {
                *client_ip = new_addr.sin_addr.s_addr;
        }
        if (client_port) {
                *client_port = new_addr.sin_port;
        }

        return new_sd;
}

int xrecv_data(sock_t sd, char * recv_buf, int recv_bytes, int timeout_ms, int timeout_ms2)
{
        struct pollfd poll_fd;
        memset(&poll_fd, 0, sizeof(poll_fd));
        poll_fd.fd = sd;

        int total_recv_bytes = 0;
        int left_bytes = recv_bytes;
        int left_time_ms = timeout_ms;
        uint32_t calc_time_start = get_cur_millseconds();
        do {
                int recv_ret = recv(sd, recv_buf+total_recv_bytes, left_bytes, 0);
                if (recv_ret > 0) {
                        total_recv_bytes += recv_ret;
                        left_bytes -= recv_ret;
                } else if (recv_ret IS 0) {
                        return -1;
                } else {
                        if (errno IS EAGAIN || errno IS EWOULDBLOCK) {
                                // continue
                        } else {
                                return recv_ret;
                        }
                }

                if (left_bytes <= 0 || left_time_ms IS 0) {
                        break;
                }

                poll_fd.events = (POLLIN | POLLRDNORM);
                poll_fd.revents = 0;
                int poll_ret = poll(&poll_fd, 1, left_time_ms);
                if (poll_ret IS -1) {
                        return -1;
                } else if (poll_ret IS 0) {
                        return total_recv_bytes;
                }
/*
                if (poll_fd.revents & (POLLERR | POLLHUP | POLLNVAL)) {
                        return -2;
                }
*/
                uint32_t calc_time_stop = get_cur_millseconds();
                if (left_time_ms < 0) {
                        if (timeout_ms2 > 0) {
                                left_time_ms = timeout_ms2;
                                calc_time_start = calc_time_stop;
                        } else {
                                left_bytes = 0;
                        }
                } else if (left_time_ms == 0) {
                        left_bytes = 0;
                } else if ((calc_time_stop - calc_time_start)>=(uint32_t)timeout_ms) {
                        left_bytes = 0;
                } else {
                        left_time_ms = timeout_ms - (calc_time_stop - calc_time_start);
                }
        } while (left_bytes > 0);

        return total_recv_bytes;

}

int xsend_data(sock_t sd, const char * send_buf, int send_bytes, int timeout_ms)
{
        struct pollfd poll_fd;
        memset(&poll_fd, 0, sizeof(poll_fd));
        poll_fd.fd = sd;

        int total_send_bytes = 0;
        int left_bytes = send_bytes;
        int left_time_ms = (timeout_ms<=0) ? 0 : timeout_ms;
        const uint32_t calc_time_start = get_cur_millseconds();
        do {
                int send_ret = send(sd, send_buf+total_send_bytes, left_bytes, MSG_NOSIGNAL);
                if (send_ret >= 0) {
                        total_send_bytes += send_ret;
                        left_bytes -= send_ret;
                } else {
                        if (send_ret IS EAGAIN || errno IS EWOULDBLOCK) {
                                // continue
                        } else {
                                return send_ret;
                        }
                }

                if (left_bytes <= 0 || left_time_ms IS 0) {
                        break;
                }

                poll_fd.events = (POLLOUT | POLLWRNORM);
                poll_fd.revents = 0;
                int poll_ret = poll(&poll_fd, 1, left_time_ms);
                if (poll_ret IS -1) {
                        return -1;
                } else if (poll_ret IS 0) {
                        return total_send_bytes;
                }
/*
                if (poll_fd.revents & (POLLERR | POLLHUP | POLLNVAL)) {
                        return -2;
                }
*/
                uint32_t calc_time_stop = get_cur_millseconds();
                if ((calc_time_stop - calc_time_start)>=(uint32_t)timeout_ms) {
                        left_bytes = 0;
                } else {
                        left_time_ms = timeout_ms - (calc_time_stop - calc_time_start);
                }
        }            while (left_bytes > 0);

        return total_send_bytes;
}

int xset_bufsize(sock_t sd, uint32_t buf_size, SNDRCV_TYPE rwt)
{
        int ret_val = -1;

        if (rwt | SOCK_SEND_TYPE) {
                ret_val = setsockopt(sd, SOL_SOCKET, SO_RCVBUF, &buf_size, sizeof(buf_size));
                if (ret_val != 0) {
                        return ret_val;
                }
        }

        if (rwt | SOCK_RECV_TYPE) {
                ret_val = setsockopt(sd, SOL_SOCKET, SO_SNDBUF, &buf_size, sizeof(buf_size));
        }

        return ret_val;
}


